from PIL import Image, ImageDraw, ImageFont
import os
from datetime import datetime
from PIL.ExifTags import TAGS, GPSTAGS
from typing import Tuple, Optional, Callable
import tkinter as tk
from tkinter import ttk, filedialog, messagebox

class ImageWatermarker:
    def __init__(self, 
                 date_format: str = '%Y-%m-%d %H:%M:%S',
                 text_color: Tuple[int, int, int] = (255, 165, 0),  # 橙色
                 font_size: int = 72,
                 font_path: Optional[str] = None,
                 margin: int = 40,
                 text_prefix: str = "拍摄时间: "):
        """
        初始化图片水印工具
        
        参数:
            date_format: 日期格式，例如：'%Y-%m-%d' 只显示日期，'%Y-%m-%d %H:%M' 显示到分钟
            text_color: RGB颜色元组，例如(255, 0, 0)为红色
            font_size: 字体大小
            font_path: 字体文件路径，如果为None则尝试使用系统字体
            margin: 文字距离边缘的距离（像素）
            text_prefix: 时间戳的前缀文本
        """
        self.date_format = date_format
        self.text_color = text_color
        self.font_size = font_size
        self.margin = margin
        self.text_prefix = text_prefix
        
        # 设置字体
        self.font = self._init_font(font_path, font_size)
        
    def _init_font(self, font_path: Optional[str], font_size: int) -> ImageFont.FreeTypeFont:
        """初始化字体"""
        try:
            if font_path:
                return ImageFont.truetype(font_path, font_size)
            
            # 尝试常用中文字体
            font_names = ["msyh.ttc", "simhei.ttf", "simsun.ttc"]
            for font_name in font_names:
                try:
                    return ImageFont.truetype(font_name, font_size)
                except:
                    continue
            
            print("警告: 未找到合适的中文字体，将使用默认字体")
            return ImageFont.load_default()
        except Exception as e:
            print(f"字体加载错误: {e}")
            return ImageFont.load_default()

    def get_capture_time(self, image_path: str) -> Optional[str]:
        """获取图片的拍摄时间"""
        try:
            with Image.open(image_path) as img:
                exif = img._getexif()
                if exif:
                    for tag_id, value in exif.items():
                        tag = TAGS.get(tag_id, tag_id)
                        if tag == 'DateTimeOriginal':
                            dt = datetime.strptime(value, '%Y:%m:%d %H:%M:%S')
                            return dt.strftime(self.date_format)
        except Exception as e:
            print(f"读取EXIF数据时出错 {image_path}: {e}")
        return None

    def add_timestamp(self, image_path: str, output_path: str) -> bool:
        """为单个图片添加时间戳水印"""
        try:
            with Image.open(image_path) as img:
                original_format = img.format
                original_size = img.size
                
                # 创建新图片
                if original_format in ['JPEG', 'JPG']:
                    new_img = Image.new('RGB', original_size, (255, 255, 255))
                    new_img.paste(img, (0, 0))
                else:
                    new_img = img.copy()
                
                draw = ImageDraw.Draw(new_img)
                
                # 获取并添加时间戳
                timestamp = self.get_capture_time(image_path)
                if timestamp:
                    text = f"{self.text_prefix}{timestamp}"
                    
                    # 计算文本位置
                    bbox = draw.textbbox((0, 0), text, font=self.font)
                    text_width = bbox[2] - bbox[0]
                    text_height = bbox[3] - bbox[1]
                    
                    x = original_size[0] - text_width - self.margin
                    y = original_size[1] - text_height - self.margin
                    
                    # 绘制文本
                    draw.text((x, y), text, font=self.font, fill=self.text_color)
                
                # 保存图片
                if original_format in ['JPEG', 'JPG']:
                    new_img.save(
                        output_path,
                        'JPEG',
                        quality=100,
                        subsampling=0,
                        optimize=False,
                        exif=img.info.get('exif', b'')
                    )
                else:
                    new_img.save(
                        output_path,
                        original_format,
                        quality=100,
                        exif=img.info.get('exif', b'')
                    )
                
                print(f"已处理: {os.path.basename(image_path)}")
                return True
                
        except Exception as e:
            print(f"处理图片时出错 {image_path}: {e}")
            import traceback
            print(traceback.format_exc())
            return False

    def process_folder(self, input_folder: Optional[str] = None, 
                      output_folder: Optional[str] = None,
                      progress_callback: Optional[Callable[[float, str], None]] = None) -> None:
        """处理文件夹中的所有图片"""
        # 设置输入输出文件夹
        input_folder = input_folder or os.getcwd()
        if output_folder is None:
            output_folder = os.path.join(input_folder, "已添加时间戳")
        
        # 创建输出文件夹
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        
        # 获取需要处理的图片列表
        image_files = [f for f in os.listdir(input_folder) 
                      if f.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp'))]
        total_files = len(image_files)
        
        if total_files == 0:
            print("\n当前文件夹中没有找到图片文件。")
            return
        
        # 处理所有图片
        for i, filename in enumerate(image_files, 1):
            input_path = os.path.join(input_folder, filename)
            output_path = os.path.join(output_folder, filename)
            self.add_timestamp(input_path, output_path)
            
            # 更新进度
            if progress_callback:
                progress = (i / total_files) * 100
                progress_callback(progress, f"正在处理: {filename} ({i}/{total_files})")
        
        print(f"\n处理完成！请在 '{output_folder}' 文件夹中查看处理后的图片。")

class WatermarkGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("图片时间戳工具")
        self.root.geometry("600x500")
        
        # 创建主框架
        main_frame = ttk.Frame(root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 日期格式选择
        ttk.Label(main_frame, text="日期格式:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.date_format_var = tk.StringVar()
        date_formats = [
            "完整格式 (2024-03-21 14:30:45)",
            "日期和时间 (2024-03-21 14:30)",
            "仅日期 (2024-03-21)",
            "中文格式 (2024年03月21日 14:30)",
            "中文日期 (2024年03月21日)"
        ]
        self.date_format_combo = ttk.Combobox(main_frame, textvariable=self.date_format_var, values=date_formats, width=30)
        self.date_format_combo.grid(row=0, column=1, sticky=tk.W, pady=5)
        self.date_format_combo.set(date_formats[0])
        
        # 文字颜色选择
        ttk.Label(main_frame, text="文字颜色:").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.color_var = tk.StringVar()
        colors = ["橙色", "红色", "绿色", "蓝色", "白色", "黑色"]
        self.color_combo = ttk.Combobox(main_frame, textvariable=self.color_var, values=colors, width=30)
        self.color_combo.grid(row=1, column=1, sticky=tk.W, pady=5)
        self.color_combo.set(colors[0])
        
        # 字体大小
        ttk.Label(main_frame, text="字体大小:").grid(row=2, column=0, sticky=tk.W, pady=5)
        self.font_size_var = tk.StringVar(value="72")
        font_size_entry = ttk.Entry(main_frame, textvariable=self.font_size_var, width=10)
        font_size_entry.grid(row=2, column=1, sticky=tk.W, pady=5)
        
        # 边距设置
        ttk.Label(main_frame, text="边距(像素):").grid(row=3, column=0, sticky=tk.W, pady=5)
        self.margin_var = tk.StringVar(value="40")
        margin_entry = ttk.Entry(main_frame, textvariable=self.margin_var, width=10)
        margin_entry.grid(row=3, column=1, sticky=tk.W, pady=5)
        
        # 文字前缀
        ttk.Label(main_frame, text="文字前缀:").grid(row=4, column=0, sticky=tk.W, pady=5)
        self.prefix_var = tk.StringVar(value="拍摄时间：")
        prefix_entry = ttk.Entry(main_frame, textvariable=self.prefix_var, width=30)
        prefix_entry.grid(row=4, column=1, sticky=tk.W, pady=5)
        
        # 输入文件夹选择
        ttk.Label(main_frame, text="输入文件夹:").grid(row=5, column=0, sticky=tk.W, pady=5)
        self.input_folder_var = tk.StringVar(value=os.getcwd())
        input_folder_entry = ttk.Entry(main_frame, textvariable=self.input_folder_var, width=50)
        input_folder_entry.grid(row=5, column=1, sticky=tk.W, pady=5)
        ttk.Button(main_frame, text="浏览", command=self.browse_input_folder).grid(row=5, column=2, padx=5)
        
        # 输出文件夹选择
        ttk.Label(main_frame, text="输出文件夹:").grid(row=6, column=0, sticky=tk.W, pady=5)
        self.output_folder_var = tk.StringVar()
        output_folder_entry = ttk.Entry(main_frame, textvariable=self.output_folder_var, width=50)
        output_folder_entry.grid(row=6, column=1, sticky=tk.W, pady=5)
        ttk.Button(main_frame, text="浏览", command=self.browse_output_folder).grid(row=6, column=2, padx=5)
        
        # 进度条
        self.progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(main_frame, length=400, mode='determinate', variable=self.progress_var)
        self.progress_bar.grid(row=7, column=0, columnspan=3, pady=20)
        
        # 开始处理按钮
        self.start_button = ttk.Button(main_frame, text="开始处理", command=self.process_images)
        self.start_button.grid(row=8, column=0, columnspan=3, pady=10)
        
        # 状态标签
        self.status_var = tk.StringVar()
        self.status_label = ttk.Label(main_frame, textvariable=self.status_var)
        self.status_label.grid(row=9, column=0, columnspan=3)

    def browse_input_folder(self):
        folder = filedialog.askdirectory(title="选择包含图片的文件夹")
        if folder:
            self.input_folder_var.set(folder)
            # 如果输出文件夹为空，自动设置默认输出路径
            if not self.output_folder_var.get():
                self.output_folder_var.set(os.path.join(folder, "已添加时间戳"))

    def browse_output_folder(self):
        folder = filedialog.askdirectory(title="选择输出文件夹")
        if folder:
            self.output_folder_var.set(folder)

    def get_color_tuple(self, color_name):
        color_map = {
            "橙色": (255, 165, 0),
            "红色": (255, 0, 0),
            "绿色": (0, 255, 0),
            "蓝色": (0, 0, 255),
            "白色": (255, 255, 255),
            "黑色": (0, 0, 0)
        }
        return color_map.get(color_name)

    def get_date_format(self, format_desc):
        format_map = {
            "完整格式 (2024-03-21 14:30:45)": '%Y-%m-%d %H:%M:%S',
            "日期和时间 (2024-03-21 14:30)": '%Y-%m-%d %H:%M',
            "仅日期 (2024-03-21)": '%Y-%m-%d',
            "中文格式 (2024年03月21日 14:30)": '%Y年%m月%d日 %H:%M',
            "中文日期 (2024年03月21日)": '%Y年%m月%d日'
        }
        return format_map.get(format_desc)

    def process_images(self):
        try:
            # 获取设置
            settings = {
                'date_format': self.get_date_format(self.date_format_var.get()),
                'text_color': self.get_color_tuple(self.color_var.get()),
                'font_size': int(self.font_size_var.get()),
                'margin': int(self.margin_var.get()),
                'text_prefix': self.prefix_var.get(),
                'input_folder': self.input_folder_var.get().strip(),  # 移除可能的空格
                'output_folder': self.output_folder_var.get().strip()  # 移除可能的空格
            }
            
            # 检查输入文件夹
            if not settings['input_folder']:
                messagebox.showerror("错误", "请选择输入文件夹！")
                return
            
            if not os.path.exists(settings['input_folder']):
                messagebox.showerror("错误", f"输入文件夹不存在：\n{settings['input_folder']}")
                return
            
            # 设置输出文件夹
            if not settings['output_folder']:
                settings['output_folder'] = os.path.join(settings['input_folder'], "已添加时间戳")
                self.output_folder_var.set(settings['output_folder'])
            
            # 确保输出文件夹的父目录存在
            output_parent = os.path.dirname(settings['output_folder'])
            if output_parent and not os.path.exists(output_parent):
                try:
                    os.makedirs(output_parent)
                except Exception as e:
                    messagebox.showerror("错误", f"无法创建输出路径：\n{output_parent}\n\n错误信息：{str(e)}")
                    return
            
            # 创建输出文件夹
            try:
                os.makedirs(settings['output_folder'], exist_ok=True)
            except Exception as e:
                messagebox.showerror("错误", f"无法创建输出文件夹：\n{settings['output_folder']}\n\n错误信息：{str(e)}")
                return
            
            # 检查输入文件夹中是否有图片
            image_files = [f for f in os.listdir(settings['input_folder']) 
                          if f.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp'))]
            if not image_files:
                messagebox.showwarning("警告", f"在输入文件夹中没有找到图片文件：\n{settings['input_folder']}")
                return
            
            # 禁用开始按钮
            self.start_button.state(['disabled'])
            self.progress_var.set(0)
            self.status_var.set("准备处理...")
            self.root.update()
            
            # 创建水印工具实例
            watermarker = ImageWatermarker(
                date_format=settings['date_format'],
                text_color=settings['text_color'],
                font_size=settings['font_size'],
                margin=settings['margin'],
                text_prefix=settings['text_prefix']
            )
            
            # 定义进度回调函数
            def update_progress(progress, status):
                self.progress_var.set(progress)
                self.status_var.set(status)
                self.root.update()
            
            # 处理图片
            watermarker.process_folder(
                input_folder=settings['input_folder'],
                output_folder=settings['output_folder'],
                progress_callback=update_progress
            )
            
            # 完成处理
            self.progress_var.set(100)
            self.status_var.set("处理完成！")
            messagebox.showinfo("完成", f"所有图片处理完成！\n\n输出文件夹：\n{settings['output_folder']}")
            
        except Exception as e:
            import traceback
            error_msg = f"处理过程中出现错误：\n{str(e)}\n\n详细信息：\n{traceback.format_exc()}"
            messagebox.showerror("错误", error_msg)
        finally:
            # 恢复开始按钮
            self.start_button.state(['!disabled'])

if __name__ == "__main__":
    root = tk.Tk()
    app = WatermarkGUI(root)
    root.mainloop()
